"""
分布式锁服务
用于在Kubernetes集群中避免定时任务重复执行
支持分片集群并行处理
"""
import logging
import uuid
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, List
import asyncio
import socket
import os

from sqlalchemy.orm import Session
from sqlalchemy import text, and_
from sqlalchemy.exc import IntegrityError

from app.database.base import get_db

logger = logging.getLogger(__name__)


class DistributedLock:
    """基于数据库的分布式锁实现"""
    
    def __init__(self, db: Session):
        self.db = db
        self.node_id = self._get_node_id()
        self.locks_held: Dict[str, str] = {}  # lock_key -> lock_token
    
    def _get_node_id(self) -> str:
        """获取节点唯一标识"""
        # 优先使用K8s Pod名称
        pod_name = os.environ.get('HOSTNAME', '')
        if pod_name:
            return pod_name
        
        # 备用方案：主机名 + 进程ID
        hostname = socket.gethostname()
        pid = os.getpid()
        return f"{hostname}-{pid}"
    
    def acquire_lock(
        self, 
        lock_key: str, 
        ttl_seconds: int = 300,
        wait_timeout: int = 0
    ) -> Optional[str]:
        """
        获取分布式锁
        
        Args:
            lock_key: 锁的唯一标识
            ttl_seconds: 锁的存活时间（秒）
            wait_timeout: 等待获取锁的超时时间（秒），0表示不等待
            
        Returns:
            str: 锁令牌，获取失败返回None
        """
        lock_token = str(uuid.uuid4())
        expire_time = datetime.now() + timedelta(seconds=ttl_seconds)
        
        start_time = datetime.now()
        
        while True:
            try:
                # 尝试插入锁记录
                result = self.db.execute(text("""
                    INSERT INTO distributed_locks (lock_key, lock_token, node_id, expire_time, created_at)
                    VALUES (:lock_key, :lock_token, :node_id, :expire_time, :created_at)
                    ON CONFLICT (lock_key) DO NOTHING
                    RETURNING lock_token
                """), {
                    'lock_key': lock_key,
                    'lock_token': lock_token,
                    'node_id': self.node_id,
                    'expire_time': expire_time,
                    'created_at': datetime.now()
                })
                
                self.db.commit()
                
                # 检查是否成功获取锁
                row = result.fetchone()
                if row:
                    self.locks_held[lock_key] = lock_token
                    logger.info(f"节点 {self.node_id} 成功获取锁 {lock_key}")
                    return lock_token
                
            except IntegrityError:
                # 锁已存在，回滚事务
                self.db.rollback()
            except Exception as e:
                logger.error(f"获取锁时发生错误: {e}")
                self.db.rollback()
                return None
            
            # 清理过期锁
            self._cleanup_expired_locks()
            
            # 检查是否超时
            if wait_timeout == 0:
                break
            
            elapsed = (datetime.now() - start_time).total_seconds()
            if elapsed >= wait_timeout:
                break
            
            # 等待一段时间后重试
            asyncio.sleep(min(1, wait_timeout - elapsed))
        
        logger.info(f"节点 {self.node_id} 获取锁 {lock_key} 失败")
        return None
    
    def release_lock(self, lock_key: str, lock_token: str) -> bool:
        """
        释放分布式锁
        
        Args:
            lock_key: 锁的唯一标识
            lock_token: 锁令牌
            
        Returns:
            bool: 释放成功返回True
        """
        try:
            result = self.db.execute(text("""
                DELETE FROM distributed_locks 
                WHERE lock_key = :lock_key AND lock_token = :lock_token
            """), {
                'lock_key': lock_key,
                'lock_token': lock_token
            })
            
            self.db.commit()
            
            if result.rowcount > 0:
                self.locks_held.pop(lock_key, None)
                logger.info(f"节点 {self.node_id} 成功释放锁 {lock_key}")
                return True
            else:
                logger.warning(f"节点 {self.node_id} 释放锁 {lock_key} 失败，锁不存在或令牌不匹配")
                return False
                
        except Exception as e:
            logger.error(f"释放锁时发生错误: {e}")
            self.db.rollback()
            return False
    
    def extend_lock(self, lock_key: str, lock_token: str, ttl_seconds: int = 300) -> bool:
        """
        延长锁的存活时间
        
        Args:
            lock_key: 锁的唯一标识
            lock_token: 锁令牌
            ttl_seconds: 延长的时间（秒）
            
        Returns:
            bool: 延长成功返回True
        """
        try:
            new_expire_time = datetime.now() + timedelta(seconds=ttl_seconds)
            
            result = self.db.execute(text("""
                UPDATE distributed_locks 
                SET expire_time = :expire_time
                WHERE lock_key = :lock_key AND lock_token = :lock_token
            """), {
                'lock_key': lock_key,
                'lock_token': lock_token,
                'expire_time': new_expire_time
            })
            
            self.db.commit()
            
            if result.rowcount > 0:
                logger.debug(f"节点 {self.node_id} 成功延长锁 {lock_key}")
                return True
            else:
                logger.warning(f"节点 {self.node_id} 延长锁 {lock_key} 失败")
                return False
                
        except Exception as e:
            logger.error(f"延长锁时发生错误: {e}")
            self.db.rollback()
            return False
    
    def _cleanup_expired_locks(self):
        """清理过期的锁"""
        try:
            result = self.db.execute(text("""
                DELETE FROM distributed_locks 
                WHERE expire_time < :current_time
            """), {
                'current_time': datetime.now()
            })
            
            self.db.commit()
            
            if result.rowcount > 0:
                logger.info(f"清理了 {result.rowcount} 个过期锁")
                
        except Exception as e:
            logger.error(f"清理过期锁时发生错误: {e}")
            self.db.rollback()
    
    def get_lock_info(self, lock_key: str) -> Optional[Dict[str, Any]]:
        """
        获取锁信息
        
        Args:
            lock_key: 锁的唯一标识
            
        Returns:
            dict: 锁信息，不存在返回None
        """
        try:
            result = self.db.execute(text("""
                SELECT lock_key, lock_token, node_id, expire_time, created_at
                FROM distributed_locks 
                WHERE lock_key = :lock_key
            """), {
                'lock_key': lock_key
            })
            
            row = result.fetchone()
            if row:
                return {
                    'lock_key': row[0],
                    'lock_token': row[1],
                    'node_id': row[2],
                    'expire_time': row[3],
                    'created_at': row[4],
                    'is_expired': row[3] < datetime.now()
                }
            
            return None
            
        except Exception as e:
            logger.error(f"获取锁信息时发生错误: {e}")
            return None
    
    def release_all_locks(self):
        """释放当前节点持有的所有锁"""
        for lock_key, lock_token in list(self.locks_held.items()):
            self.release_lock(lock_key, lock_token)
    
    def acquire_shard_lock(
        self,
        job_id: str,
        total_shards: int,
        ttl_seconds: int = 300
    ) -> Optional[int]:
        """
        获取分片锁，返回分片索引
        
        Args:
            job_id: 任务ID
            total_shards: 总分片数
            ttl_seconds: 锁的存活时间（秒）
            
        Returns:
            int: 分片索引（0-based），获取失败返回None
        """
        try:
            expire_time = datetime.now() + timedelta(seconds=ttl_seconds)
            
            # 尝试获取各个分片的锁
            for shard_index in range(total_shards):
                lock_key = f"shard:{job_id}:{shard_index}"
                lock_token = str(uuid.uuid4())
                
                try:
                    # 尝试插入分片锁记录
                    result = self.db.execute(text("""
                        INSERT INTO distributed_locks (lock_key, lock_token, node_id, expire_time, created_at, description)
                        VALUES (:lock_key, :lock_token, :node_id, :expire_time, :created_at, :description)
                        ON CONFLICT (lock_key) DO NOTHING
                        RETURNING lock_token
                    """), {
                        'lock_key': lock_key,
                        'lock_token': lock_token,
                        'node_id': self.node_id,
                        'expire_time': expire_time,
                        'created_at': datetime.now(),
                        'description': f'任务{job_id}的分片{shard_index}'
                    })
                    
                    self.db.commit()
                    
                    # 检查是否成功获取锁
                    row = result.fetchone()
                    if row:
                        self.locks_held[lock_key] = lock_token
                        logger.info(f"节点 {self.node_id} 成功获取分片锁 {lock_key}")
                        return shard_index
                        
                except IntegrityError:
                    # 该分片锁已存在，继续尝试下一个
                    self.db.rollback()
                    continue
                except Exception as e:
                    logger.error(f"获取分片锁时发生错误: {e}")
                    self.db.rollback()
                    continue
            
            # 所有分片都被占用
            logger.info(f"节点 {self.node_id} 所有分片都被占用，无法获取分片锁")
            return None
            
        except Exception as e:
            logger.error(f"获取分片锁时发生错误: {e}")
            return None
    
    def release_shard_lock(self, job_id: str, shard_index: int) -> bool:
        """
        释放分片锁
        
        Args:
            job_id: 任务ID
            shard_index: 分片索引
            
        Returns:
            bool: 释放成功返回True
        """
        lock_key = f"shard:{job_id}:{shard_index}"
        lock_token = self.locks_held.get(lock_key)
        
        if not lock_token:
            logger.warning(f"节点 {self.node_id} 未持有分片锁 {lock_key}")
            return False
        
        return self.release_lock(lock_key, lock_token)
    
    def get_active_shards(self, job_id: str) -> List[Dict[str, Any]]:
        """
        获取任务的活跃分片信息
        
        Args:
            job_id: 任务ID
            
        Returns:
            List[Dict]: 活跃分片信息列表
        """
        try:
            result = self.db.execute(text("""
                SELECT lock_key, node_id, created_at, expire_time
                FROM distributed_locks 
                WHERE lock_key LIKE :pattern
                  AND expire_time >= :current_time
                ORDER BY lock_key
            """), {
                'pattern': f'shard:{job_id}:%',
                'current_time': datetime.now()
            })
            
            shards = []
            for row in result.fetchall():
                # 解析分片索引
                lock_key = row[0]
                shard_index = int(lock_key.split(':')[-1])
                
                shards.append({
                    'shard_index': shard_index,
                    'lock_key': lock_key,
                    'node_id': row[1],
                    'created_at': row[2],
                    'expire_time': row[3]
                })
            
            return shards
            
        except Exception as e:
            logger.error(f"获取活跃分片信息时发生错误: {e}")
            return []


class LockContext:
    """分布式锁上下文管理器"""
    
    def __init__(
        self, 
        lock_key: str, 
        ttl_seconds: int = 300,
        wait_timeout: int = 0,
        db: Optional[Session] = None
    ):
        self.lock_key = lock_key
        self.ttl_seconds = ttl_seconds
        self.wait_timeout = wait_timeout
        self.db = db or next(get_db())
        self.lock = DistributedLock(self.db)
        self.lock_token: Optional[str] = None
    
    def __enter__(self):
        self.lock_token = self.lock.acquire_lock(
            self.lock_key, 
            self.ttl_seconds, 
            self.wait_timeout
        )
        
        if self.lock_token is None:
            raise Exception(f"无法获取锁: {self.lock_key}")
        
        return self.lock_token
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.lock_token:
            self.lock.release_lock(self.lock_key, self.lock_token)


# 便捷函数
def with_distributed_lock(
    lock_key: str, 
    ttl_seconds: int = 300,
    wait_timeout: int = 0
):
    """
    分布式锁装饰器
    
    Usage:
        @with_distributed_lock("task:backup", ttl_seconds=600)
        def backup_task():
            # 任务逻辑
            pass
    """
    def decorator(func):
        def wrapper(*args, **kwargs):
            with LockContext(lock_key, ttl_seconds, wait_timeout):
                return func(*args, **kwargs)
        return wrapper
    return decorator